import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from tabulate import tabulate
from colorama import Fore, Style

# 设置字体为 Times New Roman
plt.rcParams['font.family'] = 'Times New Roman'

# 读取数据
data = pd.read_excel('./assets/normalized_data.xlsx')

# 打印原始数据
print("Original Data:")
print(data)

# 删除第一列 'Years'
numeric_data = data.drop(columns=['Years'])

# 只保留数值类型的数据
numeric_data = numeric_data.apply(pd.to_numeric, errors='coerce')

# 打印数值数据
print("Numeric Data:")
print(numeric_data)

# 删除全为 NaN 的行和列
numeric_data = numeric_data.dropna(how='all', axis=0).dropna(how='all', axis=1)

# 计算相关系数矩阵
correlation_matrix = numeric_data.corr(method='pearson')

# 打印每两列之间的相关系数
print("Correlation between each pair of columns:")
correlation_pairs = []
for i in range(correlation_matrix.shape[0]):
    for j in range(i + 1, correlation_matrix.shape[0]):
        col_name_i = correlation_matrix.columns[i]  # 获取列名
        col_name_j = correlation_matrix.columns[j]  # 获取列名
        correlation = correlation_matrix.iloc[i, j]
        color = Fore.GREEN if correlation > 0 else Fore.RED
        correlation_pairs.append([col_name_i, col_name_j, f"{color}{correlation:.6f}{Style.RESET_ALL}"])

print(tabulate(correlation_pairs, headers=["Column 1", "Column 2", "Correlation"], tablefmt="grid"))

# 定义自定义颜色映射
light_pubu_colors = sns.color_palette("PuBu", 7)
light_pubu_cmap = mcolors.LinearSegmentedColormap.from_list("light_pubu_cmap", light_pubu_colors)

# 设置图形大小并绘制热力图
plt.figure(figsize=(10, 8))

# 翻译标签
translated_labels = [
    "Visitor",
    "Satisfaction",
    "Expenditure",
    "Tax",
    "Investment",
    "Melting Rate",
]

# 绘制热力图
sns.heatmap(correlation_matrix, annot=True, cmap=light_pubu_cmap, vmin=-1, vmax=1, center=0, linewidths=.5, linecolor='grey', cbar_kws={'ticks': [-0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8]}, xticklabels=translated_labels, yticklabels=translated_labels)
plt.xticks(rotation=0)
plt.yticks(rotation=0)
plt.title('Correlation Matrix Heatmap')
plt.show()